package sample.xiaosong.greedyalgorithms;

/**
 * 前缀树
 */
public class TireTree {

    static class TireNode {
        public int path;
        public int end;
        public TireNode[] next;

        public TireNode() {
            path = 0;
            end = 0;
            // 用26下标表示26个字母
            next = new TireNode[26];
        }
    }

    public TireNode root;

    public TireTree() {
        root = new TireNode();
    }

    /**
     * 将一个字符串添加到前缀树中
     *
     * 从前缀树的root节点出发,每经过一个字母,在该字母对应下标处新增一个节点(已存在不新增)
     * 每路过一个节点，该节点的path++
     * 最后一个节点的path++,end++
     * @param word 字符串
     */
    public void insert(String word) {
        if (word == null) {
            return;
        }
        char[] chars = word.toCharArray();
        int index;
        TireNode node = root;
        for (char aChar : chars) {
            // 获取字母的数组下标
            index = aChar - 'a';
            // 第一次出现的字母新增
            if (node.next[index] == null) {
                node.next[index] = new TireNode();
            }
            // 移到当前处理的字母
            node = node.next[index];
            node.path++;
        }
        node.end++;
    }

    /**
     * 将一个字符串从前缀树中移除
     *
     * 从前缀树的root节点出发,每经过一个字母,该字母对应节点的path--
     * 如果中途某个节点的path--等于0，则将该节点置位null并返回
     * 最后一个节点的path--,end--
     * @param word 字符串
     */
    public void delete(String word) {
        if (search(word) == 0) {
            return;
        }
        char[] chars = word.toCharArray();
        int index;
        TireNode node = root;
        for (char aChar : chars) {
            index = aChar - 'a';
            if (--node.next[index].path == 0) {
                node.next[index] = null;
                return;
            }
            node = node.next[index];
        }
        node.end--;
    }

    /**
     * 在前缀树中寻找字符串，并返回该字符串出现的次数
     *
     * 从树的root节点出发，一次寻找对应的字母，直到找到最后一个字母
     * 如果中途发现有后续字母下标为0，则说明前缀树不包含这个字符串，返回0
     * 返回最后一个节点的end值
     * @param word 字符串
     * @return 字符串出现的次数
     */
    public int search(String word) {
        if (word == null) {
            return 0;
        }
        char[] chars = word.toCharArray();
        int index;
        TireNode node = root;
        for (char aChar : chars) {
            index = aChar - 'a';
            if (node.next[index] == null) {
                return 0;
            }
            node = node.next[index];
        }
        return node.end;
    }

    /**
     * 在前缀树中寻找字符串，并返回该字符串作为其他字符串前缀出现的次数
     *
     * 从树的root节点出发，一次寻找对应的字母，直到找到最后一个字母
     * 如果中途发现有后续字母下标为0，则说明前缀树不包含这个字符串，返回0
     * 返回最后一个节点的path值
     * @param prefix 前缀字符串
     * @return 出现的次数
     */
    public int prefixNumber(String prefix) {
        if (prefix == null) {
            return 0;
        }
        char[] chars = prefix.toCharArray();
        int index;
        TireNode node = root;
        for (char aChar : chars) {
            index = aChar - 'a';
            if (node.next[index] == null) {
                return 0;
            }
            node = node.next[index];
        }
        return node.path;
    }

    public static void main(String[] args) {
        TireTree tree = new TireTree();
        System.out.println(tree.search("xiaosong"));
        tree.insert("xiaosong");
        System.out.println(tree.search("xiaosong"));
        tree.delete("xiaosong");
        System.out.println(tree.search("xiaosong"));
        tree.insert("xiaoaa");
        tree.insert("xiaobb");
        tree.delete("xiao");
        System.out.println(tree.search("xiao"));
        tree.delete("xiaoaa");
        System.out.println(tree.search("xiao"));
        tree.insert("xiaoqq");
        tree.insert("xiaoqw");
        tree.insert("xiaoqe");
        tree.insert("xiaoqr");
        tree.delete("xiaoqr");
        System.out.println(tree.search("xiaoqq"));
        System.out.println(tree.prefixNumber("xiaoq"));
    }
}
